# coding=utf-8
import cv2 as cv
import numpy as np
import glob
np.set_printoptions(precision=6, suppress=True)
criteria = (cv.TERM_CRITERIA_EPS + cv.TERM_CRITERIA_MAX_ITER, 30, 0.001)
fname = glob.glob("E:/Study/CSDN人工智能/1-week-vision/作业/chessboard/*.jpg")
w = 9
h = 6
objp = np.zeros((w*h,3),np.float32)
objp[:,:2] = np.mgrid[0:w,0:h].T.reshape(-1,2)
objpoints = []
imgpoints = []
for fn in fname:
    img = cv.resize(cv.imdecode(np.fromfile(fn,dtype=np.uint8),-1),dsize=(0,0),fx=0.2,fy=0.2)
    gray = cv.cvtColor(img,cv.COLOR_BGR2GRAY)
    ret, corners = cv.findChessboardCorners(gray,(w,h))
    print(ret)
    if ret:
        cv.cornerSubPix(gray,corners,(11,11),(-1,-1),criteria)
        objpoints.append(objp)
        imgpoints.append(corners)
        cv.drawChessboardCorners(img,(w,h),corners,ret)
        cv.imshow('findCorners',img)
        cv.waitKey(1)
cv.destroyAllWindows()

img = cv.resize(cv.imdecode(np.fromfile(fname[0], dtype=np.uint8), -1), dsize=(0, 0), fx=0.2, fy=0.2)
gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
ret,mtx,dist,rvecs,tvecs = cv.calibrateCamera(objpoints,imgpoints,gray.shape[::-1],None,None)
print(mtx)
print(dist)
print(rvecs[0])
print(tvecs[0])

total_error = 0
for i in range(len(objpoints)):
    imgpoints2, _ = cv.projectPoints(objpoints[i], rvecs[i], tvecs[i], mtx, dist)
    error = cv.norm(imgpoints[i],imgpoints2, cv.NORM_L2)/len(imgpoints2)
    total_error += error
print("total error: ", total_error/len(objpoints))

img = cv.resize(cv.imdecode(np.fromfile("E:/Study/CSDN人工智能/1-week-vision/作业/chessboard/pnp/a.jpg", dtype=np.uint8), -1), dsize=(0, 0), fx=0.2, fy=0.2)
gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
ret, corners = cv.findChessboardCorners(gray, (w, h))
print(cv.solvePnP(objpoints[0],corners,mtx,dist))

